#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MPI扩展性结果可视化工具

功能：
1. 读取CSV性能数据
2. 绘制扩展性曲线
3. 计算并行效率
4. 生成性能报告图表

使用方法：
    python3 plot_scaling.py results/scaling_summary_*.csv
"""

import sys
import csv
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path

# 设置中文字体支持（如果系统有中文字体）
try:
    plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
except:
    pass

def read_csv_data(csv_file):
    """读取CSV文件并提取数据"""
    cores = []
    compute_times = []
    mem_bandwidths = []
    load_imbalances = []
    parallel_effs = []
    
    with open(csv_file, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            try:
                cores.append(int(row['核数']))
                
                # 读取计算时间
                time_str = row['强扩展计算时间(s)']
                if time_str != 'N/A':
                    compute_times.append(float(time_str))
                else:
                    compute_times.append(np.nan)
                
                # 读取内存带宽
                bw_str = row['内存带宽(GB/s)']
                if bw_str != 'N/A':
                    mem_bandwidths.append(float(bw_str))
                else:
                    mem_bandwidths.append(np.nan)
                
                # 读取负载不均衡度
                imb_str = row['负载不均衡度(%)']
                if imb_str != 'N/A':
                    load_imbalances.append(float(imb_str))
                else:
                    load_imbalances.append(np.nan)
                
                # 读取并行效率
                eff_str = row['并行效率(%)']
                if eff_str != 'N/A':
                    parallel_effs.append(float(eff_str))
                else:
                    parallel_effs.append(np.nan)
            except (ValueError, KeyError) as e:
                print(f"警告: 跳过无效数据行: {row}")
                continue
    
    return {
        'cores': np.array(cores),
        'compute_times': np.array(compute_times),
        'mem_bandwidths': np.array(mem_bandwidths),
        'load_imbalances': np.array(load_imbalances),
        'parallel_effs': np.array(parallel_effs)
    }

def calculate_speedup(data):
    """计算加速比"""
    cores = data['cores']
    times = data['compute_times']
    
    # 使用第一个有效值作为基准
    valid_idx = ~np.isnan(times)
    if not np.any(valid_idx):
        return None, None
    
    base_time = times[valid_idx][0]
    base_cores = cores[valid_idx][0]
    
    # 计算实际加速比
    speedup = base_time / times
    
    # 计算理想加速比
    ideal_speedup = cores / base_cores
    
    return speedup, ideal_speedup

def plot_scaling_curves(data, output_dir):
    """绘制扩展性曲线"""
    cores = data['cores']
    times = data['compute_times']
    
    speedup, ideal_speedup = calculate_speedup(data)
    
    if speedup is None:
        print("错误: 没有有效的计算时间数据")
        return
    
    # 创建2x2子图
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    fig.suptitle('MPI并行扩展性分析', fontsize=16, fontweight='bold')
    
    # 子图1: 计算时间 vs 核数
    ax1 = axes[0, 0]
    valid_idx = ~np.isnan(times)
    ax1.plot(cores[valid_idx], times[valid_idx], 'o-', linewidth=2, markersize=8, label='实际时间')
    ax1.set_xlabel('核心数', fontsize=12)
    ax1.set_ylabel('计算时间 (秒)', fontsize=12)
    ax1.set_title('强扩展性：计算时间', fontsize=14)
    ax1.grid(True, alpha=0.3)
    ax1.legend()
    ax1.set_xscale('log', base=2)
    ax1.set_yscale('log')
    
    # 子图2: 加速比 vs 核数
    ax2 = axes[0, 1]
    valid_idx = ~np.isnan(speedup)
    ax2.plot(cores[valid_idx], speedup[valid_idx], 'o-', linewidth=2, markersize=8, label='实际加速比')
    ax2.plot(cores[valid_idx], ideal_speedup[valid_idx], '--', linewidth=2, label='理想加速比', alpha=0.7)
    ax2.set_xlabel('核心数', fontsize=12)
    ax2.set_ylabel('加速比', fontsize=12)
    ax2.set_title('加速比曲线', fontsize=14)
    ax2.grid(True, alpha=0.3)
    ax2.legend()
    ax2.set_xscale('log', base=2)
    ax2.set_yscale('log', base=2)
    
    # 子图3: 并行效率 vs 核数
    ax3 = axes[1, 0]
    efficiency = (speedup / ideal_speedup) * 100
    valid_idx = ~np.isnan(efficiency)
    ax3.plot(cores[valid_idx], efficiency[valid_idx], 's-', linewidth=2, markersize=8, color='green')
    ax3.axhline(y=100, color='r', linestyle='--', alpha=0.5, label='理想效率100%')
    ax3.axhline(y=70, color='orange', linestyle='--', alpha=0.5, label='可接受阈值70%')
    ax3.set_xlabel('核心数', fontsize=12)
    ax3.set_ylabel('并行效率 (%)', fontsize=12)
    ax3.set_title('并行效率', fontsize=14)
    ax3.grid(True, alpha=0.3)
    ax3.legend()
    ax3.set_xscale('log', base=2)
    ax3.set_ylim([0, 110])
    
    # 子图4: 内存带宽 vs 核数
    ax4 = axes[1, 1]
    mem_bw = data['mem_bandwidths']
    valid_idx = ~np.isnan(mem_bw)
    if np.any(valid_idx):
        ax4.plot(cores[valid_idx], mem_bw[valid_idx], '^-', linewidth=2, markersize=8, color='purple')
        ax4.set_xlabel('核心数', fontsize=12)
        ax4.set_ylabel('内存带宽 (GB/s)', fontsize=12)
        ax4.set_title('内存带宽扩展性', fontsize=14)
        ax4.grid(True, alpha=0.3)
        ax4.set_xscale('log', base=2)
    
    plt.tight_layout()
    
    # 保存图表
    output_file = output_dir / 'scaling_curves.png'
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"✓ 扩展性曲线已保存: {output_file}")
    
    # 显示图表（可选）
    # plt.show()

def plot_load_balance(data, output_dir):
    """绘制负载均衡分析图"""
    cores = data['cores']
    load_imb = data['load_imbalances']
    parallel_eff = data['parallel_effs']
    
    valid_imb = ~np.isnan(load_imb)
    valid_eff = ~np.isnan(parallel_eff)
    
    if not np.any(valid_imb) and not np.any(valid_eff):
        print("警告: 没有负载均衡数据")
        return
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
    fig.suptitle('负载均衡分析', fontsize=16, fontweight='bold')
    
    # 子图1: 负载不均衡度
    if np.any(valid_imb):
        ax1.bar(cores[valid_imb], load_imb[valid_imb], color='coral', alpha=0.7)
        ax1.axhline(y=5, color='green', linestyle='--', label='良好阈值5%')
        ax1.axhline(y=10, color='orange', linestyle='--', label='可接受阈值10%')
        ax1.set_xlabel('核心数', fontsize=12)
        ax1.set_ylabel('负载不均衡度 (%)', fontsize=12)
        ax1.set_title('负载不均衡度', fontsize=14)
        ax1.legend()
        ax1.grid(True, alpha=0.3, axis='y')
    
    # 子图2: 并行效率（从负载均衡测试）
    if np.any(valid_eff):
        ax2.bar(cores[valid_eff], parallel_eff[valid_eff], color='skyblue', alpha=0.7)
        ax2.axhline(y=95, color='green', linestyle='--', label='优秀阈值95%')
        ax2.axhline(y=80, color='orange', linestyle='--', label='可接受阈值80%')
        ax2.set_xlabel('核心数', fontsize=12)
        ax2.set_ylabel('并行效率 (%)', fontsize=12)
        ax2.set_title('负载均衡效率', fontsize=14)
        ax2.legend()
        ax2.grid(True, alpha=0.3, axis='y')
        ax2.set_ylim([0, 105])
    
    plt.tight_layout()
    
    output_file = output_dir / 'load_balance.png'
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"✓ 负载均衡图已保存: {output_file}")

def generate_performance_report(data, output_dir):
    """生成性能报告文本"""
    cores = data['cores']
    times = data['compute_times']
    speedup, ideal_speedup = calculate_speedup(data)
    
    if speedup is None:
        print("错误: 无法生成报告，缺少数据")
        return
    
    efficiency = (speedup / ideal_speedup) * 100
    
    report_file = output_dir / 'performance_report.txt'
    
    with open(report_file, 'w', encoding='utf-8') as f:
        f.write("=" * 80 + "\n")
        f.write("  MPI并行性能分析报告\n")
        f.write("=" * 80 + "\n\n")
        
        f.write("详细性能数据：\n")
        f.write("-" * 80 + "\n")
        f.write(f"{'核数':>8} {'时间(s)':>12} {'加速比':>12} {'理想加速比':>14} {'效率(%)':>12}\n")
        f.write("-" * 80 + "\n")
        
        valid_idx = ~np.isnan(times)
        for i in np.where(valid_idx)[0]:
            f.write(f"{cores[i]:>8} {times[i]:>12.4f} {speedup[i]:>12.2f} "
                   f"{ideal_speedup[i]:>14.2f} {efficiency[i]:>12.2f}\n")
        
        f.write("\n" + "=" * 80 + "\n")
        f.write("  性能总结\n")
        f.write("=" * 80 + "\n\n")
        
        # 找到最佳核数（效率>70%的最大核数）
        good_eff_idx = np.where((efficiency >= 70) & valid_idx)[0]
        if len(good_eff_idx) > 0:
            best_cores = cores[good_eff_idx[-1]]
            best_eff = efficiency[good_eff_idx[-1]]
            f.write(f"✓ 推荐核数: {best_cores} (并行效率: {best_eff:.1f}%)\n")
        else:
            f.write("✗ 警告: 没有核数配置达到70%的并行效率！\n")
        
        # 计算平均效率
        avg_eff = np.nanmean(efficiency[valid_idx])
        f.write(f"✓ 平均并行效率: {avg_eff:.1f}%\n")
        
        # 最大加速比
        max_speedup_idx = np.nanargmax(speedup)
        max_speedup = speedup[max_speedup_idx]
        max_speedup_cores = cores[max_speedup_idx]
        f.write(f"✓ 最大加速比: {max_speedup:.2f}x (使用{max_speedup_cores}核)\n")
        
        f.write("\n" + "=" * 80 + "\n")
        f.write("  诊断建议\n")
        f.write("=" * 80 + "\n\n")
        
        # 诊断分析
        if avg_eff < 50:
            f.write("⚠ 严重问题: 平均并行效率 < 50%\n")
            f.write("  可能原因：\n")
            f.write("  1. 算法通信开销过大\n")
            f.write("  2. 问题规模太小（强扩展极限）\n")
            f.write("  3. 硬件配置问题（NUMA、绑定）\n")
            f.write("  建议：检查通信模式、增大问题规模、优化进程绑定\n\n")
        elif avg_eff < 70:
            f.write("⚠ 一般问题: 平均并行效率在50-70%之间\n")
            f.write("  可能原因：\n")
            f.write("  1. 中等程度的通信开销\n")
            f.write("  2. 负载不够均衡\n")
            f.write("  建议：优化通信模式、检查域分解策略\n\n")
        else:
            f.write("✓ 良好: 平均并行效率 > 70%\n")
            f.write("  系统并行性能表现良好\n\n")
        
        # 检查扩展性趋势
        if len(efficiency) >= 3:
            # 计算效率下降速率
            eff_drop = efficiency[valid_idx][0] - efficiency[valid_idx][-1]
            if eff_drop > 50:
                f.write("⚠ 扩展性警告: 效率下降超过50%\n")
                f.write("  说明：强扩展性较差，不适合使用大量核心\n")
                f.write("  建议：限制最大核数，或增大问题规模（弱扩展）\n\n")
        
        f.write("=" * 80 + "\n")
    
    print(f"✓ 性能报告已保存: {report_file}")

def main():
    if len(sys.argv) < 2:
        print("使用方法: python3 plot_scaling.py <csv_file>")
        print("示例: python3 plot_scaling.py results/scaling_summary_20241103_120000.csv")
        sys.exit(1)
    
    csv_file = Path(sys.argv[1])
    
    if not csv_file.exists():
        print(f"错误: 文件不存在: {csv_file}")
        sys.exit(1)
    
    print("=" * 80)
    print("  MPI扩展性结果可视化工具")
    print("=" * 80)
    print(f"\n读取数据文件: {csv_file}")
    
    # 读取数据
    data = read_csv_data(csv_file)
    print(f"✓ 成功读取 {len(data['cores'])} 组数据")
    
    # 输出目录
    output_dir = csv_file.parent
    
    # 绘制图表
    print("\n生成图表...")
    plot_scaling_curves(data, output_dir)
    plot_load_balance(data, output_dir)
    
    # 生成报告
    print("\n生成性能报告...")
    generate_performance_report(data, output_dir)
    
    print("\n" + "=" * 80)
    print("  分析完成！")
    print("=" * 80)
    print(f"\n所有输出文件保存在: {output_dir}")
    print("  - scaling_curves.png: 扩展性曲线")
    print("  - load_balance.png: 负载均衡分析")
    print("  - performance_report.txt: 性能报告")

if __name__ == '__main__':
    main()

